from selenium import webdriver
from selenium.webdriver.chrome.service import Service
from selenium.webdriver.common.by import By
from webdriver_manager.chrome import ChromeDriverManager
import pandas as pd
import pendulum
import time
import os


def get_lottery_numbers(issue, baseurl, driver):
    url = baseurl + str(issue)
    try:
        driver.get(url)
        time.sleep(1.0)

        # 获取前区号码
        front_div = driver.find_element(By.CLASS_NAME, 'kjqQq')
        front_numbers = [num.text for num in front_div.find_elements(By.CLASS_NAME, 'kjqH')]

        # 获取后区号码
        back_div = driver.find_element(By.CLASS_NAME, 'kjqHq')
        back_numbers = [num.text for num in back_div.find_elements(By.CLASS_NAME, 'kjqL')]

        # 获取日期以及星期
        date_div = driver.find_element(By.CLASS_NAME, 'sj')
        date_spans = date_div.find_elements(By.TAG_NAME, 'span')
        draw_date = date_spans[0].text.replace('开奖时间：', '').strip()
        weekday = pendulum.parse(draw_date).day_of_week

        # 获取销售金额
        amount_div = driver.find_element(By.CLASS_NAME, 'kjqtxx')
        amount_spans = amount_div.find_elements(By.TAG_NAME, 'span')
        amount = amount_spans[1].text.replace('本期销售金额： ', '').replace('元', '').replace(',', '').strip()

        row = [issue] + [draw_date] + front_numbers + back_numbers + [weekday] + [amount]
        print(f"第{issue}期: {front_numbers} + {back_numbers} + {amount}")

        return row

    except Exception as e:
        print(f"获取第{issue}期数据失败: {str(e)}")
        return None


def main():
    # 设置Chrome浏览器选项
    options = webdriver.ChromeOptions()
    options.add_argument('--headless')
    options.add_argument('--disable-gpu')
    options.add_argument('--no-sandbox')
    options.add_argument('--disable-dev-shm-usage')

    # 初始化浏览器
    driver = webdriver.Chrome(service=Service(ChromeDriverManager().install()), options=options)

    baseurl = 'https://www.zhcw.com/kjxx/dlt/kjxq/?kjData='
    data = []  # 存储每期的开奖号码

    # 生成期号范围
    issues = []
    issues.extend(range(24126, 24153))
    issues.extend(range(25001, 25074))

    for issue in issues:
        result = get_lottery_numbers(issue, baseurl, driver)
        if result:
            data.append(result)
        time.sleep(1.0)

    # 关闭浏览器
    driver.quit()

    if data:
        columns = ['期号', '日期', '前区1', '前区2', '前区3', '前区4', '前区5', '后区1', '后区2', '星期',
                   '销售金额']
        df = pd.DataFrame(data, columns=columns)
        save_dir = "../lottery/csv"
        os.makedirs(save_dir, exist_ok=True)
        filepath = os.path.join(save_dir, "dlt_lottery_numbers.csv")
        df.to_csv(filepath, index=False, encoding='utf-8')
        print(f"\n数据已保存到: {filepath}")
        print(f"共爬取 {len(df)} 期数据")
    else:
        print("未获取到任何数据")

    if not os.path.exists("images"):
        os.makedirs("images")

if __name__ == "__main__":
    main()
